import os
import tarfile
from pathlib import Path

# Source snapshot (already in repo root)
SRC = Path(__file__).resolve().parents[2] / 'qlib_bin.tar.gz'
# Destination qlib CN daily data dir
DST = Path(os.path.expanduser('~/.qlib/qlib_data/cn_data'))
DST.mkdir(parents=True, exist_ok=True)

print('Source:', SRC)
print('Destination:', DST)

if not SRC.exists():
    raise FileNotFoundError(f'{SRC} not found. Please place qlib_bin.tar.gz in project root.')

# Extract and strip top-level components if any
with tarfile.open(SRC, 'r:gz') as tf:
    members = tf.getmembers()
    # Optionally strip first path component
    for m in members:
        parts = Path(m.name).parts
        if len(parts) > 1:
            m.name = str(Path(*parts[1:]))
        else:
            m.name = parts[0]
        tf.extract(m, path=DST)

print('Done. Data updated under', DST)